from sklearn.tree import DecisionTreeClassifier, export_graphviz
from sklearn import tree
from data_format import x_train,y_train,x_test,y_test
import numpy as np
import graphviz
import matplotlib.pyplot as plt
import pydotplus
from IPython.display import Image

# 参数
# n_classes = 3
# plot_colors = "bry"
# plot_step = 0.02

x_train = (x_train - np.min(x_train, axis=0)) / (np.max(x_train, axis=0) - np.min(x_train, axis=0))

x_test = (x_test - np.min(x_test, axis=0)) / (np.max(x_test, axis=0) - np.min(x_test, axis=0))

estimator = DecisionTreeClassifier(criterion="entropy") #  criterion用来决定不纯度的计算方法  entropy:信息熵  gini基尼系数
estimator.fit(x_train, y_train)

    # 4）模型评估
    # 方法1：直接比对真实值和预测值
y_predict = estimator.predict(x_test)
print("y_predict:\n", y_predict)
print("直接比对真实值和预测值:\n", y_test == y_predict)

    # 方法2：计算准确率
score = estimator.score(x_test, y_test)
print("准确率为：\n", score)
#属性名称，从x_train里复制的
feature_name=['LOC_TOTAL', 'LOC_BLANK', 'LOC_COMMENTS', 'LOC_CODE_AND_COMMENT',
       'LOC_EXECUTABLE', 'NUM_UNIQUE_OPERANDS', 'NUM_UNIQUE_OPERATORS',
       'NUM_OPERANDS', 'NUM_OPERATORS', 'halstead_vocabulary',
       'HALSTEAD_LENGTH', 'HALSTEAD_VOLUME', 'HALSTEAD_LEVEL',
       'HALSTEAD_DIFFICULTY', 'HALSTEAD_EFFORT', 'HALSTEAD_ERROR_EST',
       'HALSTEAD_PROG_TIME', 'BRANCH_COUNT', 'DECISION_COUNT', 'call_pairs',
       'condition_count', 'multiple_condition_count', 'CYCLOMATIC_COMPLEXITY',
       'CYCLOMATIC_DENSITY', 'DECISION_DENSITY', 'DESIGN_COMPLEXITY',
       'DESIGN_DENSITY', 'normalized_cyclomatic_complexity',
       'formal_parameters']
# 可视化决策树
dot_file='decision_tree.dot'
dot_data=tree.export_graphviz(estimator,out_file=dot_file,feature_names=feature_name,filled=True,rounded=True)
with open(dot_file) as f:
       dot_graph=f.read()

graph = pydotplus.graph_from_dot_data(dot_graph)
graph.write_pdf('decision_tree.pdf')
graph.write_png('decision_tree.png')#保存图像为pdf格式
Image(graph.create_png())   #绘制图像为png格式



#graph=graphviz.Source(dot_data)
#graph.view()

print('重要参数：\n',[*zip(feature_name,estimator.feature_importances_)])